import numpy as np
import mindspore as ms
from mindspore import nn, dataset, ops
from mindvision.engine.callback import LossMonitor

ms.context.set_context(device_target='CPU', mode=ms.PYNATIVE_MODE)

from utils import split_set



train, test = split_set()

sets = {"train": dataset.GeneratorDataset(train, ['data', 'target'], shuffle=True).batch(256),
        "test": dataset.GeneratorDataset(test, ['data', 'target']).batch(256)}


class Net(nn.Cell):
    def __init__(self, depth=4, hidden=16, auto_prefix=True, flags=None):
        super().__init__(auto_prefix, flags)
        self.depth = depth
        self.hidden = hidden

        self.rnn = nn.LSTM(input_size=1, hidden_size=self.hidden,
                           num_layers=1, batch_first=True)

        self.norm = nn.BatchNorm1d(self.depth * hidden)
        self.fc = nn.Dense(self.depth*hidden, 16)

        self.act = nn.ReLU()
        self.norm1 = nn.BatchNorm1d(16)
        self.output = nn.Dense(16, 1)

    def construct(self, *inputs):
        o, _ = self.rnn(*inputs)
        o = self.norm(o.reshape(-1, self.hidden*self.depth))
        

        o = self.act(self.fc(o))
        o = self.norm1(o)
        o = self.output(o)
        return o


net = Net(hidden=64)
loss_item = nn.MSELoss()
optim = nn.optim.SGD(net.trainable_params(), learning_rate=1e-3, momentum=0.9)

Loss = nn.WithLossCell(net, loss_item)

Eval = nn.WithEvalCell(net, loss_item,add_cast_fp32=True)
Eval.set_train(False)

Train = nn.TrainOneStepCell(Loss, optim)
Train.set_train()


if __name__ == "__main__":
    epoch = 100

    loss = nn.Loss()
    loss.clear()
    mae = nn.MAE()
    mae.clear()

    steps = sets['train'].get_batch_size()
    for i in range(epoch):
        loss_value = 0.0
        count = 0
        for it in ['train','test']:
            if it == 'train':
                for x, y in sets[it]:
                    
                    loss_value += Train(x, y)
                    if count % 10 == 9:
                        print("mode:{},epoch:[{:3d}/{:3d}],step:[{:4d}/{:4d}],loss is {:.6f}".format(
                            it, i+1, epoch, count * 10, steps, loss_value.asnumpy()/10))

                        loss_value = 0.0
                    count += 1
            else:
                mae_val = 0.0
                loss_val = 0.0
                t_val = 0.0
                count = 0
                for x, y in sets[it]:
                    outs = Eval(x, y)
                    t_val += outs[0]
                    loss.update(outs[0])
                    mae.update(outs[1],outs[2])
                    
                    if count % 10 == 9:
                        mae_val = mae.eval()
                        loss_val = loss.eval()
                        print("mode:{}: mae = {:.6f},loss = {:.6f}".format(it, mae_val, loss_val))
                        t_val = 0.0
                        loss.clear()
                        mae.clear()

                    count += 1


